import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from enum import Enum
import copy
from typing import List, Optional, Union
from math import radians
import threading
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.time import Time
from rclpy.duration import Duration
from rclpy.qos import QoSProfile
import rclpy.qos
import tf2_ros
import message_filters
from cv_bridge import CvBridge
from sensor_msgs.msg import Image, CameraInfo, PointCloud2
from utils import ObservBuffer
from navigator import Navigator
COLOR_TOPICS = ["/camera1/color/image_raw", "/camera2/color/image_raw", "/camera3/color/image_raw", "/camera4/color/image_raw"]
DEPTH_TOPICS = ["/camera1/depth/image_raw", "/camera2/depth/image_raw", "/camera3/depth/image_raw", "/camera4/depth/image_raw"]
POINTS_TOPICS = ["/camera1/depth/points", "/camera2/depth/points", "/camera3/depth/points", "/camera4/depth/points"]
COLOR_INFO_TOPICS = ["/camera1/color/camera_info", "/camera2/color/camera_info", "/camera3/color/camera_info", "/camera4/color/camera_info"]
DEPTH_INFO_TOPICS = ["/camera1/depth/camera_info", "/camera2/depth/camera_info", "/camera3/depth/camera_info", "/camera4/depth/camera_info"]


def stamp_to_sec(stamp):
    return stamp.sec + stamp.nanosec * 1e-9


class CameraParam:
    def __init__(self, frame_id: str, intrinsic_matrix: np.ndarray, resolution: Optional[tuple] = None, distortion_coeff: Optional[np.ndarray] = None):
        self.frame_id = frame_id
        self.k = intrinsic_matrix
        self.resolution = resolution
        self.distortion_coeff = distortion_coeff

    def __str__(self):
        return (
            f"Camera Name: {self.frame_id}\n"
            f"Resolution: {self.resolution[0]}x{self.resolution[1]}\n" if resolution else ""
            f"Intrinsic Matrix:\n{self.k}\n"
            f"Distortion Coefficients:\n{self.distortion_coeff}"
        )

    def __repr__(self):
        return self.__str__()

    def get_intrinsic(self):
        return self.k

    def set_intrinsic(self, camera_intrinsic):
        self.k = camera_intrinsic


class Camera(ObservBuffer):
    def __init__(self, camera_param):
        super().__init__(10)
        self.param = camera_param

    def get_intrinsic(self):
        return self.param.get_intrinsic()


class LidarCamera(ObservBuffer):
    def __init__(self):
        supper().__init__(10)
        self.cam_params: List[CameraParam] = []
        self.lidar2cam: List[np.ndarray] = []

    def add_camera(self, cam_param: CameraParam, transform: np.ndarray):
        """
        添加一个相机及其对应的变换矩阵

        :param cam_param: CameraParam 对象
        :param transform: 4x4 变换矩阵 (np.ndarray)
        """
        if transform.shape != (4, 4):
            raise ValueError("Transform matrix must be a 4x4 numpy array")
        self.cameras.append(cam_param)
        self.Tlidar2cam.append(transform)

    def __str__(self):
        details = "LidarCameraParam Details:\n"
        for i, cam in enumerate(self.cam_params):
            details += f"Camera {i}:\n{cam}\n"
        details += "Transformation Matrices from Lidar to Cameras:\n"
        for i, T in enumerate(self.Tlidar2cam):
            details += f"Tlidar2cam[{i}]:\n{T}\n"
        return details

    def __repr__(self):
        return self.__str__()


class ParamManager(Node):
    def __init__(self, node_name='param_manager', namespace=''):
        super().__init__(node_name=node_name, namespace=namespace)
        self.params = {}
        self.declare_parameter("use_point_cloud", False)
        self.declare_parameter("color_topics", COLOR_TOPICS)
        self.declare_parameter("depth_topics", DEPTH_TOPICS)
        self.declare_parameter("color_info_topics", COLOR_INFO_TOPICS)
        self.declare_parameter("depth_info_topics", DEPTH_INFO_TOPICS)
        self.declare_parameter("slop", 0.1)
        self.declare_parameter("queue_size", 10)

        self.params['use_point_cloud'] = self.get_parameter("use_point_cloud").value
        self.params['color_topics'] = self.get_parameter("color_topics").value
        self.params['depth_topics'] = self.get_parameter("depth_topics").value
        self.params['color_info_topics'] = self.get_parameter("color_info_topics").value
        self.params['depth_info_topics'] = self.get_parameter("depth_info_topics").value
        self.params['slop'] = self.get_parameter("slop").value
        self.params['queue_size'] = self.get_parameter("queue_size").value

    def get_params(self):
        return self.params

    def get_param(self, name, default=None):
        return self.params.get(name, default)

    def get_param_ros(self, name, default=None):
        self.declare_parameter(name, default)
        return self.get_parameter(name).value

    def set_params(self, params):
        self.params.update(params)

    def set_param(self, name, value):
        self.params[name] = value


class TransformBuffer(Node):
    _instance = None
    _lock = threading.Lock()

    def __new__(cls, *args, **kwargs):
        ''' singleton pattern '''
        with cls._lock:
            if not cls._instance:
                cls._instance = super(TransformBuffer, cls).__new__(cls, *args, **kwargs)
        return cls._instance

    def __init__(self, cache_duration=20.0):
        super().__init__(node_name='transform_buffer')
        self.tf_buffer = tf2_ros.Buffer(cache_time=Duration(seconds=cache_duration))
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self, spin_thread=True)

    def __del__(self):
        self.destroy_node()

    def query(self, target_frame, source_frame, time_stamp=Time(), duration=0.1):
        try:
            tf = self.tf_buffer.lookup_transform(target_frame, source_frame, time_stamp, Duration(seconds=duration))
            return tf
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            self.get_logger().error(f"Error looking up transform: {e}")
            return None

    def query_full(self, target_frame, target_time, source_frame, source_time, fixed_frame='map', duration=0.1):
        try:
            tf = self.tf_buffer.lookup_transform_full(
                target_frame=target_frame,
                target_time=target_time,
                source_frame=source_frame,
                source_time=source_time,
                fixed_frame=fixed_frame,
                timeout=Duration(seconds=duration)
            )
            return tf
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            self.get_logger().error(f"Error looking up transform: {e}")
            return None


class MultiTopicSynchronizer:
    """
    Generic class that implements time synchronization for multiple ROS topic messages.
    """
    def __init__(self, node, callback=None):
        self.node = node
        self.callback = callback

    def init(self, topics_and_msgtypes, queue_size=10, slop=None):
        """
        Initialize the MultiTopicSynchronizer.

        Parameters:
            topics_and_msgtypes (list[tuple[str, msg_type]]): A list of topics to be synchronized along with their corresponding message types.
            queue_size (int): The size of the synchronizer message queue.
            slop (float): The allowable timestamp deviation in seconds.If None, use ExactTimeSynchronizer instead of ApproximateTimeSynchronizer.
        """
        # Create a list of subscribers
        self.subscribers = []
        for topic, message_type in topics_and_msgtypes:
            qos_profile = QoSProfile(depth=queue_size, history=rclpy.qos.HistoryPolicy.KEEP_LAST,
                                     durability=rclpy.qos.DurabilityPolicy.VOLATILE,
                                     reliability=rclpy.qos.ReliabilityPolicy.BEST_EFFORT)
            subscriber = message_filters.Subscriber(self.node, message_type, topic, qos_profile=qos_profile)
            self.subscribers.append(subscriber)
        self.node.get_logger().info("The slop for TimeSynchronizer is {}".format(slop))
        if slop is None:
            self.synchronizer = message_filters.TimeSynchronizer(self.subscribers, queue_size)
        else:
            self.synchronizer = message_filters.ApproximateTimeSynchronizer(self.subscribers, queue_size, slop)
        return self

    def start(self):
        """
        Start the MultiTopicSynchronizer.
        """
        self.synchronizer.registerCallback(self.sync_callback)

    def sync_callback(self, *msgs):
        """
        Synchronized message callback function.

        Parameters:
            *msgs: A list of synchronized messages in the order of the topics list.
        """
        # Process the synchronized messages here.
        if self.callback:
            self.callback(*msgs)


class CameraInfoSubscriber:
    def __init__(self, node, topics, node_name='camera_info_subscriber', namespace=''):
        self.node = node
        self.camera_info_qos_profile = QoSProfile(depth=1, history=rclpy.qos.HistoryPolicy.KEEP_LAST,
                                                  durability=rclpy.qos.DurabilityPolicy.VOLATILE,
                                                  reliability=rclpy.qos.ReliabilityPolicy.BEST_EFFORT)
        self.num_camerainfo = len(topics)
        self.camera_infos = {}
        if isinstance(topics, dict):
            self.camera_info_topics = list(topics.values())
            self.camera_info_names = list(topics.keys())
        elif isinstance(topics, list):
            self.camera_info_topics = topics
            self.camera_info_names = [i for i in range(self.num_camerainfo)]

        self.current_info_idx = 0
        self.current_info_subscription = None
        self.lock = threading.Lock()
        self.subscribe_to_next_info()

    def subscribe_to_next_info(self):
        if self.current_info_idx < self.num_camerainfo:
            self.current_info_subscription = self.node.create_subscription(
                CameraInfo,
                self.camera_info_topics[self.current_info_idx],
                self.camera_info_callback,
                self.camera_info_qos_profile
            )
            self.node.get_logger().info(f'Waiting for camera {self.current_info_idx} info...')
        else:
            self.node.get_logger().info('All camera info received')

    def camera_info_callback(self, msg):
        self.node.get_logger().info("color info received")
        with self.lock:
            self.camera_infos[self.camera_info_names[self.current_info_idx]] = CameraParam(str(self.current_info_idx), np.array(msg.k).reshape(3, 3))
            self.node.get_logger().info("Intrinsic matrix: \n{}".format(self.camera_infos[self.camera_info_names[self.current_info_idx]].get_intrinsic()))
        # Destroy current subscription
        self.node.destroy_subscription(self.current_info_subscription)
        self.current_info_subscription = None
        # Subscribe to the next camera topic
        self.current_info_idx += 1
        self.subscribe_to_next_info()

    def len(self):
        with self.lock:
            return len(self.camera_infos)

    def get_camera_info(self, camera_name):
        with self.lock:
            if camera_name in self.camera_infos:
                return self.camera_infos[camera_name]
        return None

    def get_camera_infos(self):
        with self.lock:
            return copy.deepcopy(self.camera_infos)


class ColorDepthSyncer(MultiTopicSynchronizer):
    def __init__(self, node, parameters, process_observation=None):
        if process_observation is None:
            self.process_observation = lambda color_msgs, depth_msgs: None
        else:
            self.process_observation = process_observation

        def sync_callback(*msgs):
            self.node.get_logger().info("synchronized messages received")
            color_msgs = msgs[:self.num_color]
            depth_msgs = msgs[self.num_color:]
            self.process_observation(color_msgs, depth_msgs)

        super().__init__(node, callback=sync_callback)
        self.use_point_cloud = parameters.get('use_point_cloud', False)
        color_topics = parameters.get('color_topics', COLOR_TOPICS)
        depth_topics = parameters.get('depth_topics', DEPTH_TOPICS)
        queue_size = parameters.get('queue_size', 10)
        slop = parameters.get('slop', None)
        self.num_color = len(color_topics)
        if self.use_point_cloud:
            topics_and_msgtypes = [(topic, Image) for topic in color_topics] + [(topic, PointCloud2) for topic in depth_topics]
        else:
            topics_and_msgtypes = [(topic, Image) for topic in color_topics] + [(topic, Image) for topic in depth_topics]

        self.node.get_logger().info(f"subscribed topics and msgtypes: {topics_and_msgtypes}")
        super().init(topics_and_msgtypes, queue_size=queue_size, slop=slop)
        super().start()


class ColorDepthBuffer(ObservBuffer):
    def __init__(self, node, parameters: dict):
        max_len = parameters.get('max_len', 3)
        super().__init__(max_len=max_len)
        self.node = node
        def process_observation(color_msgs, depth_msgs):
            self.node.get_logger().info(f"processing observation")
            # 提取所有时间戳
            timestamps = [stamp_to_sec(msg.header.stamp) for msg in color_msgs + depth_msgs]
            # 计算时间戳之间的差异
            max_diff = max(timestamps) - min(timestamps)
            # time_stamp, self.observations = self.save_observation(color_msgs, depth_msgs)
            self.node.get_logger().info(f"max difference between Timestamp: {max_diff:.6f}")
            self.push((color_msgs, depth_msgs))
            self.node.get_logger().info(f"save observation {self.len()}")
        self.color_depth_syncer = ColorDepthSyncer(node, parameters, process_observation=process_observation)


class ActionType(Enum):
    STOP = 0
    MOVE_FORWARD = 1
    TURN_LEFT = 2
    TURN_RIGHT = 3


class RobotEnv:
    def __init__(self, node, parameters: dict):
        self.cv_bridge = CvBridge()
        parameters['use_point_cloud'] = False
        self.buffer = ColorDepthBuffer(node, parameters)
        self.navigator = Navigator()

    def latest_rgb(self, n=1):
        if n <= 0:
            raise ValueError("n must be a positive integer")
        observs = self.buffer.latest(n)
        if n == 1:
            observs = [observs]
        return [self.convert_observation(observ)[1][0] for observ in observs]

    def convert_observation(self, color_msgs, depth_msgs):
        if len(color_msgs) == 0:
            raise ValueError("No color messages received. At least one color message is required.")
        color_images = [self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding="rgb8") for msg in color_msgs]
        depth_images = [self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding="32FC1") for msg in depth_msgs]

        return color_msgs[0].header.stamp, color_images, depth_images

    def step(self, action: Union[ActionType, int], **kwargs) -> bool:
        """Execute one step with the given action and parameters
        
        Args:
            action: Action type to execute
        """
        if isinstance(action, int):
            action = ActionType(action)

        try:
            if action == ActionType.STOP:
                return self.navigator.cancel()
            elif action == ActionType.MOVE_FORWARD:
                distance = kwargs.get('distance', 1.0)
                if distance <= 0:
                    raise ValueError("Distance must be a positive number")
                return self.navigator.move_forward(distance)
            elif action == ActionType.TURN_LEFT:
                angle = kwargs.get('angle', 15)
                angle = radians(angle)
                return self.navigator.turn_left(angle)
            elif action == ActionType.TURN_RIGHT:
                angle = kwargs.get('angle', 15)
                angle = radians(angle)
                return self.navigator.turn_right(angle)
            else:
                self.node.get_logger().error(f"Unknown action type: {action}")
                return False
        except Exception as e:
            self.node.get_logger().error(f"Error executing action {action}: {e}")
            return False


import unittest
from unittest.mock import Mock, patch
class TestRobotEnv(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        rclpy.init()
        cls.node = rclpy.create_node('test_robot_env')
    
    @classmethod
    def tearDownClass(cls):
        cls.node.destroy_node()
        rclpy.shutdown()
    
    def setUp(self):
        self.params = {
            'color_topics': ['/test/color'],
            'depth_topics': ['/test/depth'],
            'max_len': 3
        }
        self.env = RobotEnv(self.node, self.params)
        
    def test_init(self):
        """Test initialization of RobotEnv"""
        self.assertIsNotNone(self.env.cv_bridge)
        self.assertIsNotNone(self.env.buffer)
        self.assertIsNotNone(self.env.navigator)
        
    def test_latest_rgb(self):
        """Test getting latest RGB image"""
        # Create mock color and depth messages
        color_msg = Mock(spec=Image)
        depth_msg = Mock(spec=Image)
        
        # Mock cv_bridge conversion
        mock_rgb = np.zeros((480, 640, 3), dtype=np.uint8)
        with patch.object(CvBridge, 'imgmsg_to_cv2', return_value=mock_rgb):
            self.env.buffer.push(([color_msg], [depth_msg]))
            rgb = self.env.latest_rgb(1)
            self.assertEqual(len(rgb), 1)
            self.assertEqual(rgb[0].shape, (480, 640, 3))
            
    def test_step_move_forward(self):
        """Test move forward action"""
        with patch.object(self.env.navigator, 'move_forward') as mock_move:
            mock_move.return_value = True
            result = self.env.step(ActionType.MOVE_FORWARD, distance=50)
            self.assertTrue(result)
            mock_move.assert_called_once_with(50)
            
    def test_step_turn_left(self):
        """Test turn left action"""
        with patch.object(self.env.navigator, 'turn_left') as mock_turn:
            mock_turn.return_value = True
            # Test with 90 degrees - should be converted to radians
            result = self.env.step(ActionType.TURN_LEFT, angle=90)
            self.assertTrue(result)
            mock_turn.assert_called_once()
            # Check if argument is approximately pi/2 radians
            self.assertAlmostEqual(mock_turn.call_args[0][0], 1.5708, places=4)
            
    def test_step_turn_right(self):
        """Test turn right action"""
        with patch.object(self.env.navigator, 'turn_right') as mock_turn:
            mock_turn.return_value = True
            result = self.env.step(ActionType.TURN_RIGHT, angle=45)
            self.assertTrue(result)
            mock_turn.assert_called_once()
            # Check if argument is approximately pi/4 radians
            self.assertAlmostEqual(mock_turn.call_args[0][0], 0.7854, places=4)
            
    def test_step_stop(self):
        """Test stop action"""
        with patch.object(self.env.navigator, 'cancel') as mock_cancel:
            mock_cancel.return_value = True
            result = self.env.step(ActionType.STOP)
            self.assertTrue(result)
            mock_cancel.assert_called_once()
            
    def test_invalid_action(self):
        """Test invalid action handling"""
        result = self.env.step(999)  # Invalid action
        self.assertFalse(result)

if __name__ == '__main__':
    unittest.main()